import asyncio

from langchain.globals import set_verbose
from langchain_community.chat_models import ChatTongyi
from langchain_core.messages import AIMessage, HumanMessage
from langchain_core.runnables import RunnableConfig
from langgraph.checkpoint.memory import MemorySaver
from langgraph.constants import START
from langgraph.graph import StateGraph, MessagesState
from pydantic import SecretStr
# set_verbose(True)  # 全局开启详细日志输出
chatLLM = ChatTongyi(
    model="qwen-plus-2025-04-28",   # 此处以qwen-max为例，您可按需更换模型名称。模型列表：https://help.aliyun.com/zh/model-studio/getting-started/models
    streaming=True,
    api_key = SecretStr("sk-d16b46d66abb45bb960bd9c57804e2f9"),
    # other params...
)

# r1 = chatLLM.invoke("你好，我是冬阳")
# print(r1)
# r2 = chatLLM.invoke("我叫什么名字？")
# print(r2)
#
# print("--------------------")
# result = chatLLM.invoke(
#     [
#         HumanMessage(content="你好，我是冬阳"),
#         AIMessage(content="你好，冬阳。有什么需要帮助的？"),
#         HumanMessage(content="我叫什么名字？"),
#     ]
# )
# print(result)


# Define a new graph
workflow = StateGraph(state_schema=MessagesState)


# Define the function that calls the model
async def call_model(state: MessagesState):
    response = await chatLLM.ainvoke(state["messages"])
    return {"messages": response}
# Define the (single) node in the graph
workflow.add_edge(START, "model")
workflow.add_node("model", call_model)
# Add memory
memory = MemorySaver()
app = workflow.compile(checkpointer=memory)

config:RunnableConfig = {"configurable": {"thread_id": "abc123"}}

async def chat():
    query = "你好，我是邢冬阳"

    input_messages = [HumanMessage(query)]
    output = await app.ainvoke({"messages": input_messages}, config)
    # print(output["messages"][0].pretty_print())
    # print("---------------------------")
    output["messages"][-1].pretty_print()  # output contains all messages in state

    query = "我叫什么名字?"

    input_messages = [HumanMessage(query)]
    output = await app.ainvoke({"messages": input_messages}, config)
    output["messages"][-1].pretty_print()

asyncio.run(chat())
